import asyncio
import aiofiles
from Crypto.Cipher import AES  # pip install pycryptodome
import os
import time
import requests

# 获得秘钥
def get_key(url):
    resp = requests.get(url)
    return resp.content

# 解密函数，处理一批次文件
async def aio_dec(key, start, end):
    tasks = []
    for i in range(start, end):
        task = asyncio.create_task(dec_ts(f"temp_{i}.ts", key))
        tasks.append(task)
    await asyncio.gather(*tasks)
    print(f"解密完毕: {start} 至 {end}")

# 解密单个文件
async def dec_ts(name, key):
    file_path = f"实验数据/video2/{name}"
    if not os.path.exists(file_path):
        print(f"文件不存在: {name}")
        return
    aes = AES.new(key=key, IV=b"0000000000000000", mode=AES.MODE_CBC)
    async with aiofiles.open(file_path, mode="rb") as f1, \
            aiofiles.open(f"实验数据/video2/po_{name}", mode="wb") as f2:
        bs = await f1.read()
        await f2.write(aes.decrypt(bs))

# 分批处理解密任务
async def batch_decrypt(key, num_files, max_batch_size):
    num_batches = (num_files + max_batch_size - 1) // max_batch_size
    for batch in range(num_batches):
        start = batch * max_batch_size
        end = min(start + max_batch_size, num_files)
        await aio_dec(key, start, end)

# 主函数
def main(key_url, num_files, max_batch_size=3000):
    key = get_key(key_url)
    print("从URI获取的密钥:", key)
    print("密钥长度:", len(key))

    start_time = time.time()

    asyncio.run(batch_decrypt(key, num_files, max_batch_size))
    end_time = time.time()
    print(f"用时：{end_time - start_time} 秒")

if __name__ == '__main__':
    num_files = 543  # 需要解密的文件数
    url = "https://hnzy.bfvvs.com/play/QeZLXRbp/enc.key"  # 密钥地址
    main(url, num_files)
